# 第一种融合方式: CT和PET作为双通道输入（512x512x2）+other_info

import numpy as np
import os
# import pandas as pd
# import matplotlib.pyplot as plt
# from urllib.request import urlopen, urlretrieve
# from PIL import Image
# from tqdm import tqdm_notebook
# from sklearn.utils import shuffle
# import cv2
import time
import shutil
# from sklearn.datasets import load_files
# from glob import glob
# from tensorflow.keras import applications
# from tensorflow.keras.preprocessing.image import ImageDataGenerator
# from tensorflow.keras import optimizers

from keras.models import Sequential, Model, load_model
from keras.utils import to_categorical
from keras.layers import Input, Add, Dense, Activation, ZeroPadding2D, BatchNormalization, Flatten, Conv2D
from keras.layers import AveragePooling2D, MaxPooling2D, GlobalMaxPooling2D, Dropout, MaxPool2D, GlobalAveragePooling2D
from keras.callbacks import TensorBoard, ReduceLROnPlateau, ModelCheckpoint
from keras.optimizers import  SGD, Adam
from keras.initializers import glorot_uniform


model_name = "resnet50_ct_pet_1"
time_str = time.strftime("%Y%m%d_%H%M%S", time.localtime())
batch_size = 8
log_dir = "D:/lung_cancer/lung_cancer_code/code/release_code/deepnet_output/logs/" + model_name + time_str
work_dir = "D:/lung_cancer/lung_cancer_code/code/release_code/deepnet_output/workdir/" + model_name + time_str
model_save_path = "D:/lung_cancer/lung_cancer_code/code/release_code/deepnet_output/models/" + model_name
learn_rate = 0.0001

def convert_to_one_hot(Y, C):
    Y = np.eye(C)[Y.reshape(-1)].T
    return Y

def augu_one_input(x, y):
    x_2class = [[], []]
    y_2class = [[], []]
    for i in range(len(x)):
        label = y[i]
        x_2class[label].append(x[i])
        y_2class[label].append(label)
    max_count = max([len(j) for j in x_2class])
    augu_x = []
    augu_y = []
    for label in range(2):
        n = int(max_count / len(x_2class[label]))
        augu_x.extend(x_2class[label] * n)
        augu_y.extend(y_2class[label] * n)
    return np.array(augu_x), np.array(augu_y)


def load_dataset_ct_pet_1():
    train_set_x_ct_orig = []
    train_set_x_pet_orig = []
    train_set_y_orig = []
    test_set_x_ct_orig = []
    test_set_x_pet_orig = []
    test_set_y_orig = []

    train_ct_dirs = ["D:/lung_cancer/data/Slice/two/train/CTSlice"]
    test_ct_dir = "D:/lung_cancer/data/Slice/two/test/CTSlice"

    train_pet_dirs = ["D:/lung_cancer/data/Slice/two/train/PETSlice"]
    test_pet_dir = "D:/lung_cancer/data/Slice/two/test/PETSlice"

    for i in range(len(train_ct_dirs)):
        train_ct_dir = train_ct_dirs[i]
        train_pet_dir = train_pet_dirs[i]
        for ct_file_name in os.listdir(train_ct_dir):
            ct_file_path = train_ct_dir + "/" + ct_file_name
            img_ct = np.load(ct_file_path)
            # img_ct = img_ct[:, :, np.newaxis]
            train_set_x_ct_orig.append(img_ct)

            pet_file_name = "PETSlice".join(ct_file_name.split("CTSlice"))
            pet_file_path = train_pet_dir + "/" + pet_file_name
            img_pet = np.load(pet_file_path)
            # img_pet = img_pet[:, :, np.newaxis]
            train_set_x_pet_orig.append(img_pet)

            label = int(ct_file_name.split("_")[1])
            train_set_y_orig.append(label-1)

    for ct_file_name in os.listdir(test_ct_dir):
        ct_file_path = test_ct_dir + "/" + ct_file_name
        img_ct = np.load(ct_file_path)
        # img_ct = img_ct[:, :, np.newaxis]
        test_set_x_ct_orig.append(img_ct)

        pet_file_name = "PETSlice".join(ct_file_name.split("CTSlice"))
        pet_file_path = test_pet_dir + "/" + pet_file_name
        img_pet = np.load(pet_file_path)
        # img_pet = img_pet[:, :, np.newaxis]
        test_set_x_pet_orig.append(img_pet)

        label = int(ct_file_name.split("_")[1])
        test_set_y_orig.append(label-1)

    train_set_x_ct_orig = np.array(train_set_x_ct_orig)
    train_set_x_pet_orig = np.array(train_set_x_pet_orig)
    train_set_y_orig = np.array(train_set_y_orig)
    test_set_x_ct_orig = np.array(test_set_x_ct_orig)
    test_set_x_pet_orig = np.array(test_set_x_pet_orig)
    test_set_y_orig = np.array(test_set_y_orig)

    # train test ct和pet进行合并
    train_set_x_orig = np.stack((train_set_x_ct_orig, train_set_x_pet_orig), axis=3)
    test_set_x_orig = np.stack((test_set_x_ct_orig, test_set_x_pet_orig), axis=3)


    # 平衡 不同种类的样本
    # train_set_x_orig, train_set_y_orig = augu_one_input(train_set_x_orig, train_set_y_orig)

    # shuffle trainset
    m = train_set_x_orig.shape[0]
    permutation = list(np.random.permutation(m))
    train_set_x_orig = train_set_x_orig[permutation, :, :, :]
    train_set_y_orig = train_set_y_orig[permutation]

    return train_set_x_orig, train_set_y_orig, test_set_x_orig, test_set_y_orig


def identity_block(X, f, filters, stage, block):
    """
    Implementation of the identity block as defined in Figure 3

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network

    Returns:
    X -- output of the identity block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value. You'll need this later to add back to the main path.
    X_shortcut = X

    # First component of main path
    X = Conv2D(filters=F1, kernel_size=(1, 1), strides=(1, 1), padding='valid', name=conv_name_base + '2a',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv2D(filters=F2, kernel_size=(f, f), strides=(1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv2D(filters=F3, kernel_size=(1, 1), strides=(1, 1), padding='valid', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2c')(X)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def convolutional_block(X, f, filters, stage, block, s=2):
    """
    Implementation of the convolutional block as defined in Figure 4

    Arguments:
    X -- input tensor of shape (m, n_H_prev, n_W_prev, n_C_prev)
    f -- integer, specifying the shape of the middle CONV's window for the main path
    filters -- python list of integers, defining the number of filters in the CONV layers of the main path
    stage -- integer, used to name the layers, depending on their position in the network
    block -- string/character, used to name the layers, depending on their position in the network
    s -- Integer, specifying the stride to be used

    Returns:
    X -- output of the convolutional block, tensor of shape (n_H, n_W, n_C)
    """

    # defining name basis
    conv_name_base = 'res' + str(stage) + block + '_branch'
    bn_name_base = 'bn' + str(stage) + block + '_branch'

    # Retrieve Filters
    F1, F2, F3 = filters

    # Save the input value
    X_shortcut = X

    ##### MAIN PATH #####
    # First component of main path
    X = Conv2D(F1, (1, 1), strides=(s, s), name=conv_name_base + '2a', kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2a')(X)
    X = Activation('relu')(X)

    # Second component of main path (≈3 lines)
    X = Conv2D(filters=F2, kernel_size=(f, f), strides=(1, 1), padding='same', name=conv_name_base + '2b',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2b')(X)
    X = Activation('relu')(X)

    # Third component of main path (≈2 lines)
    X = Conv2D(filters=F3, kernel_size=(1, 1), strides=(1, 1), padding='valid', name=conv_name_base + '2c',
               kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name=bn_name_base + '2c')(X)

    ##### SHORTCUT PATH #### (≈2 lines)
    X_shortcut = Conv2D(filters=F3, kernel_size=(1, 1), strides=(s, s), padding='valid', name=conv_name_base + '1',
                        kernel_initializer=glorot_uniform(seed=0))(X_shortcut)
    X_shortcut = BatchNormalization(axis=3, name=bn_name_base + '1')(X_shortcut)

    # Final step: Add shortcut value to main path, and pass it through a RELU activation (≈2 lines)
    X = Add()([X, X_shortcut])
    X = Activation('relu')(X)

    return X


def ResNet50(input_shape=(512, 512, 2), classes=1):
    """
    Implementation of the popular ResNet50 the following architecture:
    CONV2D -> BATCHNORM -> RELU -> MAXPOOL -> CONVBLOCK -> IDBLOCK*2 -> CONVBLOCK -> IDBLOCK*3
    -> CONVBLOCK -> IDBLOCK*5 -> CONVBLOCK -> IDBLOCK*2 -> AVGPOOL -> TOPLAYER

    Arguments:
    input_shape -- shape of the images of the dataset
    classes -- integer, number of classes

    Returns:
    model -- a Model() instance in Keras
    """

    # Define the input as a tensor with shape input_shape
    X_input = Input(input_shape)

    # Zero-Padding
    X = ZeroPadding2D((3, 3))(X_input)

    # Stage 1
    X = Conv2D(64, (7, 7), strides=(2, 2), name='conv1', kernel_initializer=glorot_uniform(seed=0))(X)
    X = BatchNormalization(axis=3, name='bn_conv1')(X)
    X = Activation('relu')(X)
    X = MaxPooling2D((3, 3), strides=(2, 2))(X)

    # Stage 2
    X = convolutional_block(X, f=3, filters=[64, 64, 256], stage=2, block='a', s=1)
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='b')
    X = identity_block(X, 3, [64, 64, 256], stage=2, block='c')

    ### START CODE HERE ###

    # Stage 3 (≈4 lines)
    X = convolutional_block(X, f=3, filters=[128, 128, 512], stage=3, block='a', s=2)
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='b')
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='c')
    X = identity_block(X, 3, [128, 128, 512], stage=3, block='d')

    # Stage 4 (≈6 lines)
    X = convolutional_block(X, f=3, filters=[256, 256, 1024], stage=4, block='a', s=2)
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='b')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='c')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='d')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='e')
    X = identity_block(X, 3, [256, 256, 1024], stage=4, block='f')

    # Stage 5 (≈3 lines)
    X = convolutional_block(X, f=3, filters=[512, 512, 2048], stage=5, block='a', s=2)
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='b')
    X = identity_block(X, 3, [512, 512, 2048], stage=5, block='c')

    # AVGPOOL (≈1 line). Use "X = AveragePooling2D(...)(X)"
    X = AveragePooling2D((2, 2), name="avg_pool")(X)

    ### END CODE HERE ###

    # output layer
    X = Flatten()(X)
    X = Dense(classes, activation='sigmoid', name='fc' + str(classes), kernel_initializer=glorot_uniform(seed=0))(X)

    # Create model
    model = Model(inputs=X_input, outputs=X, name=model_name)
    return model


def train():
    X_train, Y_train_orig, X_test, Y_test_orig = load_dataset_ct_pet_1()

    # Convert training and test labels to one hot matrices
    # Y_train = to_categorical(Y_train_orig)
    # Y_test = to_categorical(Y_test_orig)
    Y_train = Y_train_orig
    Y_test = Y_test_orig
    print(Y_train_orig)
    print(Y_test_orig)
    print(Y_train)
    print(Y_test)
    print("number of training examples = " + str(X_train.shape[0]))
    print("number of test examples = " + str(X_test.shape[0]))
    print("X_train shape: " + str(X_train.shape))
    print("Y_train shape: " + str(Y_train.shape))
    print("X_test shape: " + str(X_test.shape))
    print("Y_test shape: " + str(Y_test.shape))

    model = ResNet50()
    adam = Adam(lr=learn_rate)
    model.compile(optimizer=adam, loss='binary_crossentropy', metrics=['accuracy'])
    model.summary()


    if not os.path.exists(work_dir):
        os.makedirs(work_dir)
    checkpoint = ModelCheckpoint(filepath=work_dir + "/" + model_name + "_" + "_e" + "{epoch:02d}-{val_loss:.4f}.hd5",
                                 monitor='val_loss',
                                 verbose=1,
                                 save_best_only=False,
                                 save_weights_only=False,
                                 mode='auto',
                                 period=1)
    # 每隔一轮且每当val_loss降低时保存一次模型
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    checkpoint_fixed_name = ModelCheckpoint(filepath=best_model_path,
                                            monitor='val_loss',
                                            verbose=1,
                                            save_best_only=True,
                                            save_weights_only=False,
                                            mode='auto',
                                            period=1)

    model.fit(x=X_train,
              y=Y_train,
              batch_size=1,
              epochs=10,
              validation_data=(X_test, Y_test),
              callbacks=[checkpoint, checkpoint_fixed_name, TensorBoard(log_dir=log_dir)]
              )


if __name__ == "__main__":
    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    if not os.path.exists(model_save_path):
        os.makedirs(model_save_path)
    train()
    best_model_path = work_dir + "/" + model_name + "_best.hd5"
    shutil.copy(best_model_path, model_save_path)
